﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

#region

using System;
using System.Collections.Generic;
using System.Data.Services.Client;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using AzureMRCore.DataModel;
using AzureMRCore.Utils;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;

#endregion

namespace AzureMRCore.Client
{
    public class TwisterAzureClient
    {
        private static string _tableServiceAddress;
        private readonly CloudBlobClient _cloudBlobClient;
        private readonly CloudQueueClient _cloudQueueClient;
        private readonly CloudTableClient _cloudTableClient;
        private readonly JobDataContext _jobDataContext;
        private readonly string _jobID;
        private readonly CloudQueue _mapQueue;
        private readonly string _mapQueueName;
        private readonly string _mapTableName;
        private readonly CloudQueue _reduceQueue;
        private readonly string _reduceQueueName;
        private readonly CloudStorageAccount _storageAccount;

        /// <summary>
        ///   Creates an Twister4Azure client.
        /// </summary>
        /// <param name="storageAccount"> Should point to a real Azure storage account or dev storage("CloudStorageAccount.DevelopmentStorageAccount"). </param>
        /// <param name="jobID"> Unique ID for the job. </param>
        /// <param name="mapQueueName"> Azure Queue to be used for MapTask scheduling. This should match the value in the server side. </param>
        /// <param name="reduceQueueName"> Azure Queue to be used for ReduceTask scheduling. This should match the value in the server side. </param>
        public TwisterAzureClient(CloudStorageAccount storageAccount, string jobID, string mapQueueName,
                                  string reduceQueueName)
        {
            _storageAccount = storageAccount;
            _jobID = jobID;
            _mapQueueName = mapQueueName;
            _reduceQueueName = reduceQueueName;
            _tableServiceAddress = storageAccount.TableEndpoint.ToString();
            _cloudQueueClient = _storageAccount.CreateCloudQueueClient();
            _cloudTableClient = _storageAccount.CreateCloudTableClient();
            _cloudBlobClient = _storageAccount.CreateCloudBlobClient();

            if (_mapQueueName != null) _mapQueue = _cloudQueueClient.GetQueueReference(_mapQueueName);
            if (_reduceQueueName != null) _reduceQueue = _cloudQueueClient.GetQueueReference(_reduceQueueName);

            ServicePointManager.UseNagleAlgorithm = false;
            ServicePointManager.Expect100Continue = false;
            ServicePointManager.CheckCertificateRevocationList = false;
            ServicePointManager.DefaultConnectionLimit = 35;
            // ServicePointManager.SetTcpKeepAlive(true, 80000, 1000);

            _jobDataContext = CreateJobDataContext();
            _mapTableName = DataModelUtils.GetMapTableName(_jobID);
            _cloudTableClient.CreateTableIfNotExist(_mapTableName);
            _mapQueue.CreateIfNotExist();
        }

        internal TwisterAzureClient(CloudStorageAccount storageAccount, string jobID,
                                    string reduceQueueName)
        {
            _storageAccount = storageAccount;
            _jobID = jobID;
            _reduceQueueName = reduceQueueName;
            _tableServiceAddress = storageAccount.TableEndpoint.ToString();
            _cloudQueueClient = _storageAccount.CreateCloudQueueClient();
            _cloudTableClient = _storageAccount.CreateCloudTableClient();
            _cloudBlobClient = _storageAccount.CreateCloudBlobClient();

            if (_reduceQueueName != null) _reduceQueue = _cloudQueueClient.GetQueueReference(_reduceQueueName);
        }

        /* BLOB maps with iterations.  */

        public void ProcessMapRed(string mrAppName, string inputBlobContainerURI, int iteration, string programParams,
                                  int numReduceTasks, string outputContainerName,
                                  string bcastURI = null, Boolean doMerge = false)
        {
            ProcessInternal(mrAppName, null, inputBlobContainerURI, iteration, programParams, numReduceTasks,
                            outputContainerName, bcastURI, false, doMerge);
        }

        public void ProcessMapRed(string mrAppName, List<string> mapTasks, int iteration, string programParams,
                                  int numReduceTasks, string outputContainerName,
                                  string bcastURI = null, Boolean doMerge = false)
        {
            ProcessInternal(mrAppName, mapTasks, null, iteration, programParams, numReduceTasks,
                            outputContainerName, bcastURI, true, doMerge);
        }

        /// <summary>
        ///   Wait for the completion of the job specified by the job id.
        /// </summary>
        /// <param name="sleepTime"> Polling interval to check for job completion. </param>
        /// <returns> True if the job completed without errors. </returns>
        public bool WaitForCompletion(int sleepTime)
        {
            bool mapDone = false;
            var mapContext = new MapMetaDataContext(_storageAccount.TableEndpoint.ToString(),
                                                    _storageAccount.Credentials, _mapTableName);
            _cloudTableClient.CreateTableIfNotExist(mapContext.Name);

            var reduceContext = new ReduceMetaDataContext(_storageAccount.TableEndpoint.ToString(),
                                                          _storageAccount.Credentials,
                                                          DataModelUtils.GetReduceTableName(_jobID));
            _cloudTableClient.CreateTableIfNotExist(reduceContext.TableName);

            while (true)
            {
                int retry = 0;
                retryLabel:
                try
                {
                    if (!mapDone)
                    {
                        IQueryable mapTasks = mapContext.MapTasks.Where(p => p.TaskStatus != MapDataModel.FINISHED);

                        //Done processing all the map tasks
                        if (!(mapTasks.GetEnumerator().MoveNext()))
                        {
                            mapDone = true;
                            sleepTime = 1;
                        }
                    }
                    else
                    {
                        IQueryable reduceTasks =
                            reduceContext.ReduceTasks.Where(p => p.TaskStatus != MapDataModel.FINISHED);

                        //Done processing all the reduce tasks
                        if (!(reduceTasks.GetEnumerator().MoveNext()))
                        {
                            return true;
                        }
                    }
                }
                catch (Exception e)
                {
                    if (retry < 20)
                    {
                        retry++;
                        Thread.Sleep(200);
                        goto retryLabel;
                    }
                    throw new Exception("save retry failed. retry " + retry, e);
                }
                Thread.Sleep(sleepTime);
            }
        }

        public bool Cleanup(string outdir)
        {
            CloudBlobContainer intermediateDataContainer =
                _cloudBlobClient.GetContainerReference(DataModelUtils.GetIntermediateDataContainer(outdir));
            intermediateDataContainer.Delete();

            TwisterAzureClientUtils.RemoveTaskTables(_jobID, _storageAccount);
            return true;
        }

        private void ProcessInternal(string mrAppName, List<string> mapTasks, string inputDir, int iteration,
                                     string programParams, int numReduceTasks, string outputContainerName,
                                     string bcastURI,
                                     Boolean isCustom = false, Boolean doMerge = false)
        {
            // for some reason the following unused declaration seems to increase the performance. Hence leaving it.
            var mapDataContext = new MapMetaDataContext(_tableServiceAddress,
                                                        _storageAccount.Credentials,
                                                        _mapTableName);

            _mapQueue.Clear();

            var mapParams = new Dictionary<string, object>
                                {
                                    {Constants.MR_APP_NAME, mrAppName},
                                    {Constants.JOBID, _jobID},
                                    {Constants.PROGRAM_PARAMS, programParams},
                                    {Constants.NUM_REDS, numReduceTasks + ""},
                                    {Constants.OUT_DIR, outputContainerName},
                                    {Constants.ITERATION, iteration.ToString()},
                                };

            int mapTaskCount = 0;
            if (!isCustom)
            {
                mapParams.Add(Constants.TASK_TYPE, Constants.BLOB_TYPE);
                mapTaskCount = ScheduleBLOBMaps(inputDir, mapParams, bcastURI, iteration, mrAppName);
            }
            else
            {
                mapParams.Add(Constants.TASK_TYPE, Constants.CUSTOM_TYPE);
                mapTaskCount = ScheduleCustomMaps(mapTasks, mapParams, iteration, mrAppName);
            }

            if (mapTaskCount == 0)
            {
                Console.Out.WriteLine("No map tasks to process.");
                return;
            }

            var jobEntity = new JobDataModel(_jobID, numReduceTasks, programParams, outputContainerName, inputDir)
                                {
                                    PerformMerge = doMerge,
                                    NumMapTasks = mapTaskCount,
                                    Status = MapDataModel.SUBMITTED
                                };
            _jobDataContext.AddObject(_jobDataContext.Name, jobEntity);
            _jobDataContext.SaveChangesWithRetries();

            _reduceQueue.CreateIfNotExist();
            _reduceQueue.Clear();

            var reduceDataContext = new ReduceMetaDataContext(_tableServiceAddress,
                                                              _storageAccount.Credentials,
                                                              DataModelUtils.GetReduceTableName(_jobID));
            _cloudTableClient.CreateTableIfNotExist(reduceDataContext.TableName);

            var reduceParams = new Dictionary<string, object>
                                   {
                                       {Constants.MR_APP_NAME, mrAppName},
                                       {Constants.JOBID, _jobID},
                                       {Constants.PROGRAM_PARAMS, programParams},
                                       {Constants.NUM_MAPS, mapTaskCount + ""},
                                       {Constants.NUM_REDS, numReduceTasks + ""},
                                       {Constants.OUT_DIR, outputContainerName},
                                       {Constants.ITERATION, "0"}
                                   };

            ScheduleReduceTasks(reduceParams, reduceDataContext, numReduceTasks, bcastURI, mrAppName, doMerge);
        }

        private JobDataContext CreateJobDataContext()
        {
            var jobDataContext = new JobDataContext(_storageAccount.TableEndpoint.ToString(),
                                                    _storageAccount.Credentials, DataModelUtils.GetJobTableName());
            _cloudTableClient.CreateTableIfNotExist(jobDataContext.Name);

            if (jobDataContext.Jobs.Where(p => p.RowKey == _jobID).ToList().Count > 0)
            {
                Console.Out.WriteLine("JobID exists.. Please provide a different jobid.");
                throw new ArgumentException("JobID exists.. Please provide a different jobid.");
            }
            return jobDataContext;
        }

        private void ScheduleReduceTasks(IDictionary<string, object> reduceParams,
                                         ReduceMetaDataContext reduceDataContext,
                                         int numReduceTasks, string bcastURI, string mrAppName,
                                         bool doMerge = false)
        {
            for (int i = 0; i < numReduceTasks; i++)
            {
                var redTaskParams = new Dictionary<string, object>(reduceParams);
                string taskTableName = DataModelUtils.GetReduceTaskTableName(_jobID, i, mrAppName);
                redTaskParams.Add(Constants.REDUCEID, i + "_" + mrAppName);
                redTaskParams.Add(Constants.TASK_TABLE_NAME, taskTableName);
                if (i == 0)
                    AddReduceTask(reduceDataContext, redTaskParams, bcastURI, numReduceTasks, doMerge);
                else
                    AddReduceTask(reduceDataContext, redTaskParams, bcastURI, numReduceTasks);
            }
        }


        private int ScheduleCustomMaps(IList<string> mapTasks, IDictionary<string, object> mapParams, int iteration,
                                       string mrAppName)
        {
            Stopwatch stopwatch = Stopwatch.StartNew();
            var tasks = new Task[mapTasks.Count];
            //const int div = 1;
            int count = mapTasks.Count; // /div;
            int i = 0;

            mapParams.Add(Constants.NUM_MAPS, mapTasks.Count + "");
            foreach (String mapTask in mapTasks)
            {
                var mapTaskParams = new Dictionary<string, object>(mapParams)
                                        {
                                            {Constants.MAPID, i},
                                            {Constants.CUSTOM_KEY, mapTask},
                                            {Constants.CUSTOM_VALUE, mapTask}
                                        };
                int i1 = i;
                tasks[i] = new Task(() => AddCustomMapTask(mapTaskParams, i1, iteration, mrAppName, count));
                i++;
            }
            for (int j = 0; j < mapTasks.Count; j++)
            {
                tasks[j].Start();
            }

            stopwatch.Stop();
            Console.Out.WriteLine("Time for schedule:" + mapTasks.Count + "  " + stopwatch.ElapsedMilliseconds);
            return mapTasks.Count;
        }

        private int ScheduleBLOBMaps(string inputDir, IDictionary<string, object> mapParams, string bcastURI,
                                     int iteration, string mrAppName)
        {
            CloudBlobContainer container = _cloudBlobClient.GetContainerReference(inputDir);
            int count = 0;

            int i = 0;
            var options = new BlobRequestOptions
                              {
                                  BlobListingDetails = BlobListingDetails.None
                              };
            IEnumerable<IListBlobItem> list = container.ListBlobs(options);
            count = list.Count();

            var tasks = new Task[count];
            mapParams.Add(Constants.NUM_MAPS, count + "");

            string partialMessage = JSONUtils.ToPartialJson(mapParams);

            foreach (IListBlobItem item in container.ListBlobs(options))
            {
                string itemUri = item.Uri.ToString();
                string mapID = i + "_" + mrAppName;

                tasks[i] =
                    new Task(() => AddBlobMapTask(mapID, iteration, mrAppName, itemUri, bcastURI, partialMessage, count));
                tasks[i].Start();
                //tasks[i] = new Task(() => AddMapTaskToQueue(partialMessage, itemUri, mapID));
                //tasks[i].Start();

                i++;
            }

            //Block until all tasks complete.
            Task.WaitAll(tasks);

            return count;
        }

        private void AddCustomMapTask(Dictionary<string, object> mapParams, int id, int iteration, string app,
                                      int numMapTasks)
        {
            string message = JSONUtils.ToJson(mapParams);
            AddMapTask(id.ToString(), null, message, iteration, app, null, numMapTasks);
        }

        private void AddBlobMapTask(string mapID, int iteration, string appName, string blobURL, string bcastDataURL,
                                    string partialMessage, int numMapTasks)
        {
            string message = JSONUtils.ToJsonFromPartial(partialMessage, mapID, blobURL);
            AddMapTask(mapID, blobURL, message, iteration, appName, bcastDataURL, numMapTasks);
        }

        private void AddMapTask(string mapID, string blobURL, string message, int iteration, string appName,
                                string bcastDataURL, int numMapTasks)
        {
            _mapQueue.BeginAddMessage(new CloudQueueMessage(message), null, null);

            var mapDataContext = new MapMetaDataContext(_storageAccount.TableEndpoint.ToString(),
                                                        _storageAccount.Credentials,
                                                        _mapTableName);
            // mapDataContext.MergeOption = MergeOption.NoTracking;
            //Trace.WriteLine((string) mapParams[Constants.MAPID]);
            var mapTaskEntity = new MapDataModel(mapID, iteration.ToString(),
                                                 appName,
                                                 blobURL)
                                    {
                                        Attempt1Status = MapDataModel.SUBMITTED,
                                        Attempt2Status = MapDataModel.PENDING,
                                        Attempt3Status = MapDataModel.PENDING,
                                        TaskStatus = MapDataModel.SUBMITTED,
                                        BroadcastData = bcastDataURL,
                                        NumMapTasks = numMapTasks
                                    };

            mapDataContext.AddObject(mapDataContext.Name, mapTaskEntity);
            mapDataContext.BeginSaveChangesWithRetries(SaveChangesOptions.Batch, null, new object());
        }

        /*  private void AddMapTaskToQueue(string partialMessage, string blobURL, string mapId)
        {
            if (_mapQueue != null)
            {
                string message = JSONUtils.ToJsonFromPartial(partialMessage, mapId, blobURL);
                _mapQueue.AddMessage(new CloudQueueMessage(message));
            }
        }*/


        private void AddReduceTask(ReduceMetaDataContext reduceDataContext, Dictionary<string, object> reduceParams,
                                   string bcastDataURL, int numReds,
                                   bool doMerge = false)
        {
            var reduceTaskEntity = new ReduceDataModel((string) reduceParams[Constants.REDUCEID],
                                                       (string) reduceParams[Constants.ITERATION],
                                                       (string) reduceParams[Constants.MR_APP_NAME])
                                       {
                                           Attempt1Status = MapDataModel.SUBMITTED,
                                           Attempt2Status = MapDataModel.PENDING,
                                           Attempt3Status = MapDataModel.PENDING,
                                           TaskStatus = MapDataModel.SUBMITTED,
                                           BroadcastData = bcastDataURL,
                                           NumReduceTasks = numReds,
                                           RunMerge = doMerge
                                       };
            reduceDataContext.AddObject(reduceDataContext.TableName, reduceTaskEntity);
            reduceDataContext.SaveChangesWithRetries();

            string message = JSONUtils.ToJson(reduceParams);
            _reduceQueue.AddMessage(new CloudQueueMessage(message));
        }

        #region iterative

/*        // Schedule iterative tasks.. To be used internally by merge tasks...
        internal int AddIterativeBLOBMaps(JobDataModel jobModel, string mrAppName, int iteration, string bcastURI,
                                          MapMetaDataContext mapDataContext)
        {
            var mapParams = new Dictionary<string, object>
                                {
                                    {Constants.MR_APP_NAME, mrAppName},
                                    {Constants.JOBID, _jobID},
                                    {Constants.PROGRAM_PARAMS, jobModel.ProgramParams},
                                    {Constants.NUM_REDS, jobModel.NumReduceTasks + ""},
                                    {Constants.OUT_DIR, jobModel.OutDir},
                                    {Constants.ITERATION, iteration.ToString()},
                                    {Constants.TASK_TYPE, Constants.BLOB_TYPE}
                                };
            return ScheduleBLOBMaps(jobModel.InDir, mapParams, bcastURI, iteration, mrAppName);
        }*/

        internal void AddIterativeReduceTasks(JobDataModel jobModel, string mrAppName, int iteration, string bcastURI,
                                              int numReduceTasks, ReduceMetaDataContext reduceDataContext, bool isMerge)
        {
            var reduceParams = new Dictionary<string, object>
                                   {
                                       {Constants.MR_APP_NAME, mrAppName},
                                       {Constants.JOBID, _jobID},
                                       {Constants.PROGRAM_PARAMS, jobModel.ProgramParams},
                                       {Constants.NUM_MAPS, jobModel.NumMapTasks + ""},
                                       {Constants.NUM_REDS, numReduceTasks.ToString()},
                                       {Constants.ITERATION, iteration.ToString()},
                                       {Constants.OUT_DIR, jobModel.OutDir}
                                   };

            ScheduleReduceTasks(reduceParams, reduceDataContext, numReduceTasks, bcastURI, mrAppName,
                                isMerge);
        }

        #endregion
    }
}